# 演示如何利用pinocchio联合mujoco进行impedance控制
# 注意两边数据的存取与转换
# mujoco读取末端接触力以及方向，并且读取关节力矩
# 根据受力大小和方向，以及设定的目标电位，计算出末端期望速度以及方向
# 映射到关节空间


import pinocchio as pin
import mujoco
import mujoco.viewer as viewer
import numpy as np
from os.path import dirname, join
import time
from numpy.linalg import pinv
import matplotlib.pyplot as plt

# 在viewer循环前初始化存储
vel_diff_history = []

# pinocchio加载模型
modelpath = join(dirname(__file__), 'Dof6Arm.xml')
pin_model = pin.RobotWrapper.BuildFromMJCF(modelpath)

# mujoco加载模型
model = mujoco.MjModel.from_xml_path(join(dirname(__file__),'scene.xml')) 
data = mujoco.MjData(model)

key_id = model.key("home").id
mujoco.mj_resetDataKeyframe(model, data, key_id)

print(f"qpos : {data.qpos}")
print(f"type of qpos : {type(data.qpos)}")

damping: float = 1e-4
diag = damping * np.eye(6)
qa_max = 1.0
site_id = model.site('site_Hand').id
base_id = model.body('BigArm1').id
goal_pos = [0.5, 0.5, 0.5]
end_effector_id = model.body('Hand').id
jac = np.zeros((6, model.nv))

mocap_id = model.body("target").mocapid[0]

# pinocchio和mujoco都要更新
pin.framesForwardKinematics(pin_model.model, pin_model.data, data.qpos)
pin.updateFramePlacements(pin_model.model, pin_model.data) 
mujoco.mj_step(model, data)            # step xdot= f(x,u)
mujoco.mj_forward(model, data)

pos_sensor_names = [
        "Motor_Shoulder_pos", 
        "Motor_BigArm1_pos",
        "Motor_BigArm2_pos",
        "Motor_SmallArm1_pos",
        "Motor_SmallArm2_pos",
        "Motor_Hand_pos"
        ]

vel_sensor_names = [
        "Motor_Shoulder_vel", 
        "Motor_BigArm1_vel",
        "Motor_BigArm2_vel",
        "Motor_SmallArm1_vel",
        "Motor_SmallArm2_vel",
        "Motor_Hand_vel"
        ]

torque_sensor_names = [
    "Motor_Shoulder_Torque",
    "Motor_BigArm1_Torque", 
    "Motor_BigArm2_Torque",
    "Motor_SmallArm1_Torque",
    "Motor_SmallArm2_Torque"
]

# 获取目标位置SE3矩阵T_t
mocap_pos = data.mocap_pos[mocap_id]
mocap_quat = data.mocap_quat[mocap_id]

# 创建Pinocchio四元数对象（参数顺序为x,y,z,w）
pin_quat = pin.Quaternion(mocap_quat[1], mocap_quat[2], mocap_quat[3], mocap_quat[0])
# 转换为旋转矩阵
mocap_rot = pin_quat.toRotationMatrix()

global T_t0
T_t0 = pin.SE3(mocap_rot, mocap_pos)




# TODO (暂时不好做)考虑mocap的速度，把速度误差纳入到IK里面
# TODO 考虑动力补偿和积分误差
# TODO 最最重要一点。把求解IK和控制分开！！！

def ik(model, data):
    
    # 通过sensor获取关节角度比直接使用data更符合实际
    qpos_sensors = np.array([data.sensor(pos_sensor_name).data[0] for pos_sensor_name in pos_sensor_names])
    qvel_sensors = np.array([data.sensor(vel_sensor_name).data[0] for vel_sensor_name in vel_sensor_names])

    pin.framesForwardKinematics(pin_model.model, pin_model.data, qpos_sensors)
    pin.updateFramePlacements(pin_model.model, pin_model.data)


    if not hasattr(ik, 'prev_v_star'):  # 初始化前次控制量
        ik.prev_v_star = qvel_sensors

    # 在ik函数中添加完整动力学计算
    tau_gravity = pin.computeGeneralizedGravity(pin_model.model, pin_model.data, qpos_sensors)

    # tau_coriolis = pin.computeCoriolisMatrix(pin_model.model, pin_model.data, qpos_sensors, qvel_sensors) @ qvel_sensors

    # 前馈补偿项 = 重力 + 科氏力 + 离心力
    # tau_diff = tau_coriolis

    tau_measured = np.array([data.sensor(name).data[0] for name in torque_sensor_names])
    tau_measured = np.append(tau_measured, 0.0)  # 补充末端关节，假设第6关节无传感器


    tau_diff = tau_measured - tau_gravity

    k_tau = -1e-1

    tau_comp = k_tau * tau_diff

    # 获取目标位置SE3矩阵T_t
    mocap_pos = data.mocap_pos[mocap_id]
    mocap_quat = data.mocap_quat[mocap_id]
    
    # 创建Pinocchio四元数对象（参数顺序为x,y,z,w）
    pin_quat = pin.Quaternion(mocap_quat[1], mocap_quat[2], mocap_quat[3], mocap_quat[0])
    # 转换为旋转矩阵
    mocap_rot = pin_quat.toRotationMatrix()
    T_t = pin.SE3(mocap_rot, mocap_pos)
    
    # # 求出target位姿差值，用来计算target移动速度
    # T_tv = T_t.actInv(T_t0)
    # ev_q = pin.log(T_tv).vector

    # 求当前site的位姿T_b和目标位姿的相对差值矩阵T_bt
    T_bt = pin_model.data.oMf[pin_model.model.getFrameId("site_Hand")].actInv(T_t)
    # 定义的误差向量
    e_q = pin.log(T_bt).vector


    # 求site_Hand的雅克比J_b
    J_b = pin.computeFrameJacobian(
        pin_model.model, 
        pin_model.data, 
        qpos_sensors, 
        pin_model.model.getFrameId("site_Hand"), 
        pin.ReferenceFrame.LOCAL_WORLD_ALIGNED  # 根据需求选择坐标系
    )


    
    J_l = -pin.Jlog6(T_bt.inverse())
    # 接下来会这里会利用J_error v^{*} = -e(q)来求v^{*}
    J_error = J_l.dot(J_b) # 误差Jacobian
    # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    v_star = np.linalg.solve(J_error.T.dot(J_error) + diag, -J_error.T.dot(e_q))
    # v_star = pinv(J_error).dot(-e_q) # 此处是求解的关键，这里直接用pinv，没有借用damping项
    # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
    # 速度控制时候使用，TODO 存在抖动，可能是没有把误差速度加入到状态空间
    # ctrl = v_star + tau_comp
    # 将速度严格限制在[-1, 1]区间
    v_star = np.clip(v_star, -1.0, 1.0)

    # 计算允许变化量
    delta_v = v_star - ik.prev_v_star
    delta_v_clipped = np.clip(delta_v, -qa_max*model.opt.timestep, qa_max*model.opt.timestep)
    v_star = ik.prev_v_star + delta_v_clipped
    ik.prev_v_star = v_star.copy()

    ctrl = v_star + tau_comp

    return ctrl




with viewer.launch_passive(model,data) as viewer:
    # Initialize the camera view to that of the free camera.
    mujoco.mjv_defaultFreeCamera(model, viewer.cam)
    step = 0
    time_array = []
    while viewer.is_running():

        data.ctrl = ik(model,data)  # apply control
        mujoco.mj_step(model, data)            # step xdot= f(x,u)
        mujoco.mj_forward(model, data)
        qvel_sensors = np.array([data.sensor(vel_sensor_name).data[0] for vel_sensor_name in vel_sensor_names])
        vel_diff = data.qvel - data.ctrl
        vel_diff_history.append(vel_diff.copy())  # 记录历史数据
        step += 1
        t_now = model.opt.timestep * step
        time_array.append(t_now)
        
        print(f"vel_diff = {vel_diff}")
        viewer.sync()
        time.sleep(model.opt.timestep)
    
    print(f"type of vel_diff_history: {type(vel_diff_history)}")
    vel_diff_array = np.array(vel_diff_history)
    for j in range(vel_diff_array.shape[1]):
        plt.plot(time_array, vel_diff_array[:,j])
    
    plt.title('Joint Velocity Tracking Errors')
    plt.xlabel('Time (s)')
    plt.ylabel('Error (rad/s)')
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.show()